package main

import (
    "database/sql"
    "log"
    "net/http"
    "time"
    "websocket-master"
    _ "mysql-master"
)

// MySQL配置
const (
    dbUser     = "root"
    dbPassword = "haohao"
    dbName     = "mydatabase"
)

// WebSocket配置
var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

var db *sql.DB

func initDB() {
    var err error
    dsn := dbUser + ":" + dbPassword + "@/" + dbName
    db, err = sql.Open("mysql", dsn)
    if err != nil {
        log.Fatalf("Error opening database: %v", err)
    }
    if err = db.Ping(); err != nil {
        log.Fatalf("Error connecting to database: %v", err)
    }
}

func queryData() ([]map[string]interface{}, error) {
    rows, err := db.Query("SELECT * FROM student") // 替换为您的查询语句
    if err != nil {
        return nil, err
    }
    defer rows.Close()

    columns, err := rows.Columns()
    if err != nil {
        return nil, err
    }

    var results []map[string]interface{}
    for rows.Next() {
        values := make([]interface{}, len(columns))
        valuePtrs := make([]interface{}, len(columns))
        for i := range values {
            valuePtrs[i] = &values[i]
        }

        if err := rows.Scan(valuePtrs...); err != nil {
            return nil, err
        }

        result := make(map[string]interface{})
        for i, col := range columns {
            var v interface{}
            val := values[i]
            b, ok := val.([]byte)
            if ok {
                v = string(b)
            } else {
                v = val
            }
            result[col] = v
        }
        results = append(results, result)
    }
    return results, nil
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("Failed to upgrade to websocket:", err)
        return
    }
    defer conn.Close()

    for {
        // 定时查询数据并发送给客户端
        data, err := queryData()
        if err != nil {
            log.Println("Error querying data:", err)
            return
        }

        err = conn.WriteJSON(data)
        if err != nil {
            log.Println("Error sending data:", err)
            return
        }

        time.Sleep(5 * time.Second) // 每5秒查询一次数据
    }
}

func main() {
    initDB()
    defer db.Close()

    http.HandleFunc("/ws", handleWebSocket)
    log.Println("WebSocket server listening on :8081")
    log.Fatal(http.ListenAndServe(":8081", nil))
}
